use crate::printer::ASTPrinter;
use crate::token::{TextSpan, Token};

pub struct Ast {
  statements: Vec<ASTStatement>,
}

impl Ast {
  pub fn new() -> Self {
    Self {
      statements: Vec::new(),
    }
  }

  pub fn add_statement(&mut self, statement: ASTStatement) {
    self.statements.push(statement)
  }

  pub fn visit(&mut self, visitor: &mut dyn ASTVisitor) {
    for statement in &self.statements {
      visitor.visit_statement(statement);
    }
  }

  pub fn visualize(&mut self) {
    let mut printer = ASTPrinter::new();
    self.visit(&mut printer);

    println!("{}", printer.result);
  }
}

pub enum ASTStatementKind {
  Expression(ASTExpression),
}

pub struct ASTStatement {
  pub kind: ASTStatementKind,
}

impl ASTStatement {
  pub fn new(kind: ASTStatementKind) -> Self {
    Self { kind }
  }

  pub fn expression(expr: ASTExpression) -> Self {
    Self::new(ASTStatementKind::Expression(expr))
  }
}

pub enum ASTExpressionKind {
  Number(ASTNumberExpression),               // 34
  Binary(ASTBinaryExpression),               // 4 + 2 * 3 + 1
  Parenthesized(ASTParenthesizedExpression), // (4 + 2) * 3 + 1
  Unary(ASTUnaryExpression),                 // - 负号，! 非 一元运算符
  Boolean(ASTBooleanExpression),             // 布尔值 true/false
  Error(TextSpan),
}

pub struct ASTBooleanExpression {
  pub value: bool,
  pub token: Token,
}

#[derive(Debug)]
pub enum ASTUnaryOperatorKind {
  Minus,      // - 负号运算符
  BitwiseNot, // ! 非运算符
}

pub struct ASTUnaryOperator {
  pub kind: ASTUnaryOperatorKind,
  pub token: Token,
}

impl ASTUnaryOperator {
  pub fn new(kind: ASTUnaryOperatorKind, token: Token) -> Self {
    Self { kind, token }
  }
}

pub struct ASTUnaryExpression {
  pub operator: ASTUnaryOperator,  // 负号运算符
  pub operand: Box<ASTExpression>, // 操作数
}

#[derive(Debug)]
pub enum ASTBinaryOperatorKind {
  Plus,     // 加法
  Minus,    // 减法
  Multiply, // 乘法
  Divide,   // 除法
}

pub struct ASTBinaryOperator {
  pub kind: ASTBinaryOperatorKind,
  pub token: Token,
}

impl ASTBinaryOperator {
  pub fn new(kind: ASTBinaryOperatorKind, token: Token) -> Self {
    Self { kind, token }
  }

  pub fn precedence(&self) -> u8 {
    match self.kind {
      ASTBinaryOperatorKind::Plus | ASTBinaryOperatorKind::Minus => 1,
      ASTBinaryOperatorKind::Multiply | ASTBinaryOperatorKind::Divide => 2,
    }
  }
}

pub struct ASTNumberExpression {
  pub number: f64,
}

pub struct ASTBinaryExpression {
  pub lhs: Box<ASTExpression>,
  pub operator: ASTBinaryOperator,
  pub rhs: Box<ASTExpression>,
}

pub struct ASTParenthesizedExpression {
  pub expression: Box<ASTExpression>,
}

pub struct ASTExpression {
  pub kind: ASTExpressionKind,
}

impl ASTExpression {
  pub fn new(kind: ASTExpressionKind) -> Self {
    Self { kind }
  }

  // 34
  pub fn number(number: f64) -> Self {
    Self::new(ASTExpressionKind::Number(ASTNumberExpression { number }))
  }

  // 4 + 2 * 3
  pub fn binary(operator: ASTBinaryOperator, lhs: Self, rhs: Self) -> Self {
    Self::new(ASTExpressionKind::Binary(ASTBinaryExpression {
      lhs: Box::new(lhs),
      operator,
      rhs: Box::new(rhs),
    }))
  }

  // (4 + 2) * 3
  pub fn parenthesized(expression: Self) -> Self {
    Self::new(ASTExpressionKind::Parenthesized(
      ASTParenthesizedExpression {
        expression: Box::new(expression),
      },
    ))
  }

  // -  !
  pub fn unary(operator: ASTUnaryOperator, operand: ASTExpression) -> Self {
    ASTExpression::new(ASTExpressionKind::Unary(ASTUnaryExpression {
      operator,
      operand: Box::new(operand),
    }))
  }

  // true false
  pub fn boolean(token: Token, value: bool) -> Self {
    ASTExpression::new(ASTExpressionKind::Boolean(ASTBooleanExpression {
      token,
      value,
    }))
  }

  pub fn error(span: TextSpan) -> Self {
    ASTExpression::new(ASTExpressionKind::Error(span))
  }
}

pub trait ASTVisitor {
  fn do_visit_statement(&mut self, statement: &ASTStatement) {
    match &statement.kind {
      ASTStatementKind::Expression(expr) => {
        self.visit_expression(expr);
      }
    }
  }

  fn do_visit_expression(&mut self, expression: &ASTExpression) {
    match &expression.kind {
      ASTExpressionKind::Number(number) => self.visit_number_expression(number),
      ASTExpressionKind::Binary(expr) => self.visit_binary_expression(expr),
      ASTExpressionKind::Parenthesized(expr) => self.visit_parenthesized_expression(expr),
      ASTExpressionKind::Unary(expr) => self.visit_unary_expression(expr),
      ASTExpressionKind::Boolean(expr) => self.visit_boolean_expression(expr),
      ASTExpressionKind::Error(span) => self.visit_error(span),
    }
  }

  fn visit_statement(&mut self, statement: &ASTStatement) {
    self.do_visit_statement(statement);
  }

  fn visit_expression(&mut self, expression: &ASTExpression) {
    self.do_visit_expression(expression);
  }

  fn visit_number_expression(&mut self, number: &ASTNumberExpression);

  fn visit_binary_expression(&mut self, binary_expression: &ASTBinaryExpression) {
    self.visit_expression(&binary_expression.lhs);
    self.visit_expression(&binary_expression.rhs);
  }

  fn visit_parenthesized_expression(
    &mut self,
    parenthesized_expression: &ASTParenthesizedExpression,
  ) {
    self.visit_expression(&parenthesized_expression.expression);
  }

  fn visit_unary_expression(&mut self, unary_expression: &ASTUnaryExpression);

  fn visit_boolean_expression(&mut self, boolean: &ASTBooleanExpression);

  fn visit_error(&mut self, span: &TextSpan);
}
